from .base import MetricBase
from torchmetrics.image import StructuralSimilarityIndexMeasure

class SSIM(MetricBase):
    def __init__(self, device, **kwargs):
        super().__init__()
        self.metric = StructuralSimilarityIndexMeasure(data_range=(0,1), **kwargs).to(device)
        
    def compute(self, y_true, y_pred):
        return self.metric(y_true, y_pred)
        
    def summary(self):
        ssim = self.value / self.total
        print(f'SSIM: {ssim}')
        return ssim
    
if __name__ == "__main__":
    ssim = SSIM()
    import torch
    
    for _ in range(5):
        a = torch.randint(0, 2, size=(16,1,64,64)).float()
        b = torch.randint(0, 2, size=(16,1,64,64)).float()
        ssim(a, b)
        ssim.summary()
        print(ssim.total)
    ssim.summary()